class ModelEnsemble:
    def __init__(self, models):
        self.models = models
        
    def predict(self, X):
        """集成预测"""
        predictions = []
        for model in self.models:
            pred = model.predict(X)
            predictions.append(pred)
        
        # 平均预测概率
        avg_pred = np.mean(predictions, axis=0)
        return avg_pred
    
    def evaluate_ensemble(self, X_test, y_test):
        """评估集成模型"""
        y_pred_proba = self.predict(X_test)
        y_pred = np.argmax(y_pred_proba, axis=1)
        y_true = np.argmax(y_test, axis=1)
        
        from sklearn.metrics import classification_report, confusion_matrix
        print("Classification Report:")
        print(classification_report(y_true, y_pred))
        
        plt.figure(figsize=(10, 8))
        sns.heatmap(confusion_matrix(y_true, y_pred), annot=True, fmt='d', cmap='Blues')
        plt.title('Confusion Matrix')
        plt.show()

# 创建集成模型
cnn_model = model_builder.build_cnn_model()
resnet_model = model_builder.build_resnet_model()
transformer_model = model_builder.build_transformer_model()

models = [cnn_model, resnet_model, transformer_model]
ensemble = ModelEnsemble(models)

# 训练所有模型
for i, model in enumerate(models):
    print(f"Training model {i+1}")
    trainer = ModelTrainer(model, f'model_{i+1}')
    trainer.compile_model()
    trainer.train_model(X_train.reshape(-1, X_train.shape[1], 1), y_train_cat, 
                       X_test.reshape(-1, X_test.shape[1], 1), y_test_cat,
                       epochs=50, batch_size=32)

# 评估集成模型
ensemble.evaluate_ensemble(X_test.reshape(-1, X_test.shape[1], 1), y_test_cat)
